#include <types.h>
#include <arm32.h>
#include <kernel/initcall.h>


static uint32_t mmu_l1_tbl[4096] align(0x4000) = {0};


typedef enum {
    MMU_SECTION_NCNB = 0x0, // not cached, not buffered
    MMU_SECTION_NCB  = 0x1, // not cached, buffered
    MMU_SECTION_CNB  = 0x2, // cached, writethrough
    MMU_SECTION_CB   = 0x3, // cached, writeback
} mmu_section_type_e;


/**
 * @brief  map a section in the L1 page table.
 * @param table Pointer to the L1 page table.
 * @param virt Virtual address to map.
 * @param phys Physical address to map.
 * @param size Size of the section in bytes.
 * @param type Type of the section.
 * @return None
 * @note This function maps a section in the L1 page table, allowing the MMU to translate virtual addresses to physical addresses.
 */
static inline void mmu_map_l1_entry( uint32_t* table,
                                           uint32_t virt, uint32_t phys,
                                           uint32_t size, mmu_section_type_e type)
{
    /* L1 virtual address entry, 1MB  */
    virt >>= 20;
    /* L1 physical address entry, 1MB  */
    phys >>= 20;
    /* size in MB */
    size >>= 20;
    type &= 0x3;

    for(volatile uint32_t i = 0; i < size; i++) {
        table[virt] = (phys << 20) | (0x3 << 10) | (0x0 << 5) | (type << 2) | (0x2 << 0);
        virt++;
        phys++;
    }
}


static int mmu_init(void)
{
    // Initialize translation table
    mmu_map_l1_entry(mmu_l1_tbl, 0x00000000, 0x00000000, SZ_2G, MMU_SECTION_NCNB);
    mmu_map_l1_entry(mmu_l1_tbl, 0x80000000, 0x80000000, SZ_2G, MMU_SECTION_NCNB);
    // Enable caching for first 64MB of DRAM
    mmu_map_l1_entry(mmu_l1_tbl, 0x80000000, 0x80000000, SZ_1M * 64, MMU_SECTION_CB);

    /*  Setup MMU TLB address */
    arm32_ttb_set((uint32_t)(mmu_l1_tbl));
    /* Invalidate MMU TLB */
    arm32_tlb_invalidate();
    /* Set domain access control */
    // Domain 0: Client, Domain 1: Manager, Domain 2-15: No access
    arm32_domain_set(0x3);
    /* Enable MMU  */
    arm32_mmu_enable();
    /* Enable instruction cache */
    arm32_icache_enable();
    /* Enable data cache */
    arm32_dcache_enable();

    return 0;
}


early_initcall(mmu_init);
